"""
ETL流程功能演示示例

该示例展示了使用SQLServerManager类实现完整的ETL（抽取-转换-加载）流程，
包括数据抽取、清洗转换、加载到目标表以及全程事务管理和错误处理。
"""
from modules.database.sqlserver_manager import SQLServerManager
from modules.dataframe.dataframe import Table
import logging
import pandas as pd
from datetime import datetime

# 配置日志
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
logger = logging.getLogger('sqlserver_manager')
logger.addHandler(handler)
logger.setLevel(logging.INFO)

def transform_data(source_data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """
    数据转换函数：清洗和转换原始数据
    """
    transformed = []
    
    for row in source_data:
        # 1. 数据清洗：处理缺失值
        if not row.get('Name') or not row.get('Department'):
            logger.warning(f"跳过不完整记录: {row}")
            continue
        
        # 2. 数据转换：标准化部门名称
        dept_mapping = {
            '工程': 'Engineering',
            '市场': 'Marketing',
            '销售': 'Sales',
            '人力资源': 'HR',
            '财务': 'Finance'
        }
        normalized_dept = dept_mapping.get(row['Department'], row['Department'])
        
        # 3. 数据计算：计算年薪和入职年限
        try:
            monthly_salary = float(row['Salary'])
            annual_salary = round(monthly_salary * 12, 2)
            
            hire_date = datetime.strptime(row['HireDate'], '%Y-%m-%d')
            years_of_service = datetime.now().year - hire_date.year
            
            # 4. 创建转换后的记录
            transformed_row = {
                'EmployeeID': row['ID'],
                'FullName': row['Name'].strip(),
                'Department': normalized_dept,
                'MonthlySalary': monthly_salary,
                'AnnualSalary': annual_salary,
                'HireDate': hire_date,
                'YearsOfService': years_of_service,
                'ETLProcessedDate': datetime.now()
            }
            transformed.append(transformed_row)
        except (ValueError, TypeError) as e:
            logger.error(f"数据转换失败: {str(e)}, 记录: {row}")
            continue
    
    logger.info(f"数据转换完成: {len(transformed)}/{len(source_data)} 条记录成功转换")
    return transformed

def demonstrate_etl_process():
    # 数据库连接参数
    db_config = {
        'server': 'localhost',
        'database': 'test_db',
        'username': 'sa',
        'password': 'your_password',
        'port': 1433
    }

    # 使用上下文管理器确保连接正确关闭
    with SQLServerManager(**db_config) as db:
        try:
            source_table = 'ETLSourceEmployees'
            target_table = 'ETLTargetEmployees'

            # 1. 准备ETL环境
            logger.info("=== 准备ETL环境 ===")
            # 清理可能存在的旧表
            for table in [source_table, target_table]:
                if db.is_table_exists(table):
                    db.execute_non_query(f"DROP TABLE {table}")
            
            # 创建源表
            logger.info(f"创建源表 {source_table}...")
            db.execute_non_query(f"""
                CREATE TABLE {source_table} (
                    ID INT PRIMARY KEY IDENTITY(1,1),
                    Name NVARCHAR(50),
                    Department NVARCHAR(50),
                    Salary NVARCHAR(20),
                    HireDate NVARCHAR(20)
                )
            """)
            
            # 创建目标表
            logger.info(f"创建目标表 {target_table}...")
            db.execute_non_query(f"""
                CREATE TABLE {target_table} (
                    EmployeeID INT PRIMARY KEY,
                    FullName NVARCHAR(50) NOT NULL,
                    Department NVARCHAR(50) NOT NULL,
                    MonthlySalary DECIMAL(10,2) NOT NULL,
                    AnnualSalary DECIMAL(10,2) NOT NULL,
                    HireDate DATE NOT NULL,
                    YearsOfService INT NOT NULL,
                    ETLProcessedDate DATETIME NOT NULL
                )
            """)
            
            # 插入源数据（包含一些不规范数据用于演示清洗过程）
            source_data = [
                {'Name': '张三', 'Department': '工程', 'Salary': '8000', 'HireDate': '2018-03-15'},
                {'Name': '李四 ', 'Department': '市场', 'Salary': '7500', 'HireDate': '2020-07-22'},
                {'Name': '王五', 'Department': '销售', 'Salary': '9000', 'HireDate': '2019-01-10'},
                {'Name': '赵六', 'Department': '财务', 'Salary': '8500', 'HireDate': '2017-09-05'},
                {'Name': '', 'Department': 'HR', 'Salary': '7000', 'HireDate': '2021-05-30'},  # 空名称
                {'Name': '钱七', 'Department': 'IT', 'Salary': 'abc', 'HireDate': '2022-01-15'},   # 薪资非数字
                {'Name': '孙八', 'Department': '人力资源', 'Salary': '9500', 'HireDate': '2016-04-20'}
            ]
            db.batch_insert(source_table, source_data)
            logger.info(f"插入 {len(source_data)} 条原始数据到源表")

            # 2. 开始ETL事务
            logger.info("\n=== 开始ETL流程 ===")
            db.begin_transaction()

            # 3. 数据抽取（Extract）
            logger.info("抽取源数据...")
            extracted_data = db.execute_query(f"SELECT * FROM {source_table}")
            logger.info(f"成功抽取 {len(extracted_data)} 条记录")

            # 4. 数据转换（Transform）
            logger.info("转换数据...")
            transformed_data = transform_data(extracted_data)

            # 使用DataFrame进行数据分析和转换验证
            logger.info("使用DataFrame验证转换结果...")
            df = Table(transformed_data)
            logger.info(f"转换后数据统计信息:\n{df.describe()}")

            # 5. 数据加载（Load）
            logger.info("加载数据到目标表...")
            if transformed_data:
                affected_rows = db.bulk_insert(target_table, transformed_data)
                logger.info(f"成功加载 {affected_rows} 条记录到目标表")
            else:
                logger.warning("没有可加载的转换后数据")

            # 6. 验证加载结果
            logger.info("验证加载结果...")
            target_count = db.execute_query(f"SELECT COUNT(*) AS Count FROM {target_table}")[0]['Count']
            if target_count == len(transformed_data):
                logger.info(f"数据验证成功: 目标表包含 {target_count} 条记录")
                db.commit()
                logger.info("ETL事务已提交")
            else:
                raise Exception(f"数据验证失败: 目标表记录数({target_count})与转换后记录数({len(transformed_data)})不匹配")

        except Exception as e:
            db.rollback()
            logger.error(f"ETL流程失败: {str(e)}", exc_info=True)
        finally:
            # 7. 清理演示环境
            logger.info("\n=== 清理演示环境 ===")
            db.execute_non_query(f"DROP TABLE IF EXISTS {source_table}")
            db.execute_non_query(f"DROP TABLE IF EXISTS {target_table}")
            logger.info("ETL演示环境清理完成")

if __name__ == "__main__":
    demonstrate_etl_process()